#!/bin/bash
# 标题: SM4 目录文件并发加密/解密性能测试脚本
# 描述: 接收输入目录和最大并发数。自动检测所有文件，执行 NUMA 绑定并发加解密，并将加密文件输出到指定目录。

# --- 配置项 ---
TEST_BIN="./hct_speed" # 您的加密测试工具路径
CMD_TYPE="sw_avx" # 默认命令类型
TEMP_OUTPUT_DIR="/tmp/sm4_test_results" # 临时日志存放目录

# 文件大小/位计算 (不再需要固定值，但保留变量以便后续使用)
# FILE_SIZE_BYTES=$((FILE_SIZE_MB * 1024 * 1024))
# FILE_SIZE_BITS=$(awk "BEGIN {print $FILE_SIZE_BYTES * 8}")

# --- 检查依赖和初始化 ---

# 检查测试工具是否存在
if [ ! -x "$TEST_BIN" ]; then
    echo "错误: 找不到测试程序 $TEST_BIN 或没有执行权限。" >&2
    echo "请检查 TEST_BIN 变量配置是否正确。" >&2
    exit 1
fi

# 检查 numactl 是否存在
if ! command -v numactl &> /dev/null; then
    echo "错误: numactl 命令未安装。请安装 numactl 以支持 NUMA 绑核。" >&2
    exit 1
fi

# --- 参数解析和目录设置 ---

# 显示用法信息
show_usage() {
    echo "使用方法: $0 <模式> <参数> <最大并发数> [选项]" >&2
    echo "" >&2
    echo "模式 A (内存中创建测试文件):" >&2
    echo "  $0 -m <文件大小> <文件数量> <最大并发数> [-cmd auto|ccp|sw_avx] [-n]" >&2
    echo "  示例: $0 -m 1MB 100 64" >&2
    echo "  示例: $0 -m 1MB 100 64 -cmd ccp" >&2
    echo "  示例: $0 -m 1MB 100 64 -n  (不输出加密文件，显示解密性能)" >&2
    echo "" >&2
    echo "模式 B (指定目录或文件列表):" >&2
    echo "  $0 -d <输入目录> <最大并发数> [-cmd auto|ccp|sw_avx] [-n]" >&2
    echo "  示例: $0 -d /path/to/data 64" >&2
    echo "  示例: $0 -d /path/to/data 64 -cmd auto" >&2
    echo "  示例: $0 -d /path/to/data 64 -n  (不输出加密文件，显示解密性能)" >&2
    echo "" >&2
    echo "文件大小格式: 支持 KB, MB, GB (如 1MB, 10KB, 2GB)" >&2
    echo "选项:" >&2
    echo "  -cmd auto|ccp|sw_avx  指定命令类型 (默认: sw_avx)" >&2
    echo "  -n  不输出加密文件，显示加密和解密性能数据" >&2
}

# 解析文件大小字符串为字节数
parse_file_size() {
    local size_str="$1"
    local size_num=$(echo "$size_str" | sed 's/[^0-9]//g')
    local size_unit=$(echo "$size_str" | sed 's/[0-9]//g' | tr '[:lower:]' '[:upper:]')
    
    case "$size_unit" in
        "KB"|"K")
            echo $((size_num * 1024))
            ;;
        "MB"|"M")
            echo $((size_num * 1024 * 1024))
            ;;
        "GB"|"G")
            echo $((size_num * 1024 * 1024 * 1024))
            ;;
        *)
            # 默认认为是字节
            echo "$size_num"
            ;;
    esac
}

# 初始化变量
OUTPUT_ENCRYPTED_FILES=true
MODE=""
FILE_SIZE_STR=""
FILE_COUNT=""
INPUT_DIR=""
MAX_CONCURRENCY=""
TEST_MODE=""

# 解析参数
while [ $# -gt 0 ]; do
    case "$1" in
        -m|--memory)
            TEST_MODE="memory"
            if [ $# -lt 4 ]; then
                echo "错误: 模式 A 需要 3 个参数: 文件大小、文件数量、最大并发数" >&2
                show_usage
                exit 1
            fi
            FILE_SIZE_STR="$2"
            FILE_COUNT="$3"
            MAX_CONCURRENCY="$4"
            shift 4
            
            # 验证文件数量
            if ! [[ "$FILE_COUNT" =~ ^[0-9]+$ ]] || [ "$FILE_COUNT" -lt 1 ]; then
                echo "错误: 文件数量必须是一个大于等于 1 的整数。" >&2
                exit 1
            fi
            ;;
        -d|--directory)
            TEST_MODE="directory"
            if [ $# -lt 3 ]; then
                echo "错误: 模式 B 需要 2 个参数: 输入目录、最大并发数" >&2
                show_usage
                exit 1
            fi
            INPUT_DIR="$2"
            MAX_CONCURRENCY="$3"
            shift 3
            
            if [ ! -d "$INPUT_DIR" ]; then
                echo "错误: 输入目录 $INPUT_DIR 不存在或不是一个目录。" >&2
                exit 1
            fi
            ;;
        -cmd)
            if [ $# -lt 2 ]; then
                echo "错误: -cmd 参数需要一个值 (auto|ccp|sw_avx)" >&2
                show_usage
                exit 1
            fi
            CMD_TYPE="$2"
            case "$CMD_TYPE" in
                auto|ccp|sw_avx)
                    ;;
                *)
                    echo "错误: 无效的命令类型 '$CMD_TYPE'，必须是 auto|ccp|sw_avx" >&2
                    show_usage
                    exit 1
                    ;;
            esac
            shift 2
            ;;
        -n)
            OUTPUT_ENCRYPTED_FILES=false
            shift
            ;;
        *)
            echo "错误: 未知参数 '$1'" >&2
            show_usage
            exit 1
            ;;
    esac
done

# 检查必需参数
if [ -z "$TEST_MODE" ]; then
    echo "错误: 必须指定测试模式 (-m 或 -d)" >&2
    show_usage
    exit 1
fi

if [ -z "$MAX_CONCURRENCY" ]; then
    echo "错误: 必须指定最大并发数" >&2
    show_usage
    exit 1
fi

# 验证最大并发数
if ! [[ "$MAX_CONCURRENCY" =~ ^[0-9]+$ ]] || [ "$MAX_CONCURRENCY" -lt 1 ]; then
    echo "错误: 最大并发数必须是一个大于等于 1 的整数。" >&2
    exit 1
fi

# 定义输出目录
if [ "$TEST_MODE" = "directory" ]; then
    OUTPUT_DIR="${INPUT_DIR}_sm4_out"
else
    OUTPUT_DIR="/tmp/sm4_memory_test_out"
fi

# 根据模式设置文件列表
if [ "$TEST_MODE" = "directory" ]; then
    # 模式 B: 从目录获取文件列表
    INPUT_FILES=($(find "$INPUT_DIR" -type f))
    N_FILES=${#INPUT_FILES[@]}
    
    if [ "$N_FILES" -eq 0 ]; then
        echo "错误: 在目录 $INPUT_DIR 中未找到任何文件进行测试。" >&2
        exit 1
    fi
    
    echo "--- 文件系统并发测试配置 ---" >&2
    echo "输入目录: $INPUT_DIR" >&2
    echo "输出目录: $OUTPUT_DIR" >&2
    echo "检测到文件总数: $N_FILES" >&2
    echo "最大并发数限制: $MAX_CONCURRENCY" >&2
    echo "命令类型: $CMD_TYPE" >&2
    echo "--------------------------------" >&2
else
    # 模式 A: 在临时目录中创建测试文件
    N_FILES="$FILE_COUNT"
    FILE_SIZE_BYTES=$(parse_file_size "$FILE_SIZE_STR")
    MEMORY_DIR="/tmp/sm4_benchmark_$$"  # 使用进程ID避免冲突
    
    echo "--- 内存模式配置 ---" >&2
    echo "文件大小: $FILE_SIZE_STR ($FILE_SIZE_BYTES 字节)" >&2
    echo "文件数量: $N_FILES" >&2
    echo "临时目录: $MEMORY_DIR" >&2
    echo "输出目录: $OUTPUT_DIR" >&2
    echo "最大并发数限制: $MAX_CONCURRENCY" >&2
    echo "命令类型: $CMD_TYPE" >&2
    echo "--------------------------------" >&2
    
    # 创建临时目录
    mkdir -p "$MEMORY_DIR"
    
    # 并发创建测试文件
    echo "正在并发创建 $N_FILES 个测试文件，最大并发数: $MAX_CONCURRENCY..." >&2
    file_creation_pids=()
    file_creation_counter=0
    
    for i in $(seq 1 $N_FILES); do
        INPUT_FILES+=("$MEMORY_DIR/testfile_$i.dat")
        
        # 并发创建文件
        dd if=/dev/urandom of="${INPUT_FILES[-1]}" bs="$FILE_SIZE_BYTES" count=1 2>/dev/null &
        pid=$!
        file_creation_pids+=($pid)
        file_creation_counter=$((file_creation_counter + 1))
        
        # 检查并发限制
        if [ "${#file_creation_pids[@]}" -ge "$MAX_CONCURRENCY" ]; then
            wait -n
            # 更新正在运行的进程数
            file_creation_pids=($(jobs -p))
        fi
    done
    
    # 等待所有文件创建完成
    for pid in "${file_creation_pids[@]}"; do
        wait "$pid" 2>/dev/null || echo "警告: 文件创建进程 $pid 失败" >&2
    done
    echo "测试文件创建完成" >&2
    
    # 预热功能已移除
    echo "跳过预热阶段..." >&2
fi

# --- NUMA 信息收集 ---

declare -A NODE_CORES
N_NODES=0
ALL_CORE_IDS=""

echo "--- 正在收集系统 NUMA 信息 ---"
NUMA_HARDWARE_OUTPUT=$(numactl --hardware)

N_NODES=$(echo "$NUMA_HARDWARE_OUTPUT" | grep 'available' | awk '{print $2}')
if [ -z "$N_NODES" ] || [ "$N_NODES" -eq 0 ]; then
    echo "警告: 无法确定 NUMA 节点数，将使用单节点模式 (N_NODES=1)。" >&2
    N_NODES=1
fi

# 功能: 展开核心ID范围 (例如 "0-3, 5" 转换为 "0 1 2 3 5")
expand_core_list() {
    local raw_list="$1"
    local expanded_list=""
    
    local cleaned_list=$(echo "$raw_list" | sed 's/[^0-9,-]/ /g' | tr ',' ' ' | xargs)

    for item in $cleaned_list; do
        if [[ "$item" =~ ^[0-9]+-[0-9]+$ ]]; then
            local start=$(echo "$item" | cut -d '-' -f 1)
            local end=$(echo "$item" | cut -d '-' -f 2)
            expanded_list+="$(seq $start $end) "
        elif [[ "$item" =~ ^[0-9]+$ ]]; then
            expanded_list+="$item "
        fi
    done
    echo "$expanded_list" | xargs
}

for i in $(seq 0 $((N_NODES - 1))); do
    RAW_CORE_LIST=$(echo "$NUMA_HARDWARE_OUTPUT" | grep "node $i cpus:" | cut -d ':' -f 2 | xargs)
    CORE_LIST=$(expand_core_list "$RAW_CORE_LIST")
    
    NODE_CORES[$i]="$CORE_LIST"
    ALL_CORE_IDS="$ALL_CORE_IDS $CORE_LIST"
done

ALL_CORE_IDS=$(echo "$ALL_CORE_IDS" | xargs -n1 | sort -n | uniq | tr '\n' ' ')
N_CORES=$(echo "$ALL_CORE_IDS" | wc -w)

echo "检测到 NUMA 节点数: $N_NODES"
echo "检测到总 CPU 核心数: $N_CORES"
echo "--------------------------------"

# 功能: 根据任务ID获取绑定的NUMA节点（均分到所有NUMA节点）
get_numa_node_for_job() {
    local job_id=$1
    local total_jobs=$2
    
    # 简单轮询分配：将任务均分到所有NUMA节点
    local node_id=$((job_id % N_NODES))
    echo "$node_id"
}

# 创建输出目录和临时日志目录，并设置清理
mkdir -p "$OUTPUT_DIR"
mkdir -p "$TEMP_OUTPUT_DIR"
rm -rf "$TEMP_OUTPUT_DIR"/* # 清理日志目录
trap "rm -rf $TEMP_OUTPUT_DIR $MEMORY_DIR" EXIT




# 功能: 运行文件系统并发测试 (已修改以处理文件列表和并发限制)
# 参数: None (使用全局变量 $INPUT_FILES, $MAX_CONCURRENCY, $OUTPUT_DIR)
# 输出: (stdout) 总耗时 (秒)
run_concurrent_test() {
    local n="$N_FILES"
    local MAX_P="$MAX_CONCURRENCY"
    local PIDS=()
    local JOB_MAP=()
    local START_TIME=$(date +%s%N)
    local JOB_COUNTER=0
    local CURRENT_JOBS=0
    
    # 构建测试环境变量
    local TEST_EVN="-elapsed -engine hct -cmd $CMD_TYPE -evp sm4"

    #echo "[Jobs] 正在启动 $n 个文件加密测试，最大并发限制为 $MAX_P..." >&2

    for INPUT_FILE in "${INPUT_FILES[@]}"; do
        
        # 1. 确定输出路径和日志路径
        if [ "$TEST_MODE" = "directory" ]; then
            # 模式 B: 从输入文件路径构建相对路径
            local RELATIVE_PATH="${INPUT_FILE#$INPUT_DIR/}"
            local OUTPUT_SUBDIR=$(dirname "$RELATIVE_PATH")
            local OUTPUT_BASE=$(basename "$RELATIVE_PATH")
            local FINAL_OUTPUT_DIR="$OUTPUT_DIR/$OUTPUT_SUBDIR"
            local OUTPUT_FILE="$FINAL_OUTPUT_DIR/$OUTPUT_BASE.enc"
        else
            # 模式 A: 直接使用文件名，不构建子目录结构
            local OUTPUT_BASE=$(basename "$INPUT_FILE")
            local FINAL_OUTPUT_DIR="$OUTPUT_DIR"
            local OUTPUT_FILE="$FINAL_OUTPUT_DIR/$OUTPUT_BASE.enc"
        fi
        
        local OUTPUT_LOG="$TEMP_OUTPUT_DIR/result_$JOB_COUNTER.log"

        mkdir -p "$FINAL_OUTPUT_DIR" # 确保输出目录存在

        # 2. 确定 NUMA 绑定（均分到所有NUMA节点）
        local numa_node=$(get_numa_node_for_job "$JOB_COUNTER" "$n")

        # 3. 运行命令
        # 根据是否输出加密文件决定是否使用 -outfile 参数
        if [ "$OUTPUT_ENCRYPTED_FILES" = "true" ]; then
            # 输出加密文件，只显示加密性能
            numactl -N $numa_node -l \
                "$TEST_BIN" $TEST_EVN -infile "$INPUT_FILE" -outfile "$OUTPUT_FILE" \
                > "$OUTPUT_LOG" 2>&1 &
        else
            # 不输出加密文件，显示加密和解密性能
            numactl -N $numa_node -l \
                "$TEST_BIN" $TEST_EVN -infile "$INPUT_FILE" \
                > "$OUTPUT_LOG" 2>&1 &
        fi
        
        local PID=$!
        
        PIDS+=($PID)
        JOB_MAP+=("$PID:$JOB_COUNTER:$INPUT_FILE")
        
        CURRENT_JOBS=$((CURRENT_JOBS + 1))
        JOB_COUNTER=$((JOB_COUNTER + 1))
        
        # JOB_COUNTER 在此行用于显示，是 1-based index
        #echo "  [Job $JOB_COUNTER/$n] 绑定核心: $CORE_ID, 内存节点: $MEM_NODE, 文件: $INPUT_FILE" >&2
        
        # 4. 检查并发限制并等待
        if [ "$CURRENT_JOBS" -ge "$MAX_P" ]; then
            # 等待任一作业完成，释放一个插槽
            wait -n
            # 更新当前正在运行的作业计数
            CURRENT_JOBS=$(jobs -r | wc -l)
        fi
    done

    # 5. 等待所有剩余的作业
    for pid in "${PIDS[@]}"; do
        wait "$pid" 2>/dev/null
    done
    local END_TIME=$(date +%s%N)

    # 6. 检查日志并统计成功数
    local SUCCESS=0
    for job_entry in "${JOB_MAP[@]}"; do
        local pid=$(echo "$job_entry" | cut -d ':' -f 1)
        local job_i=$(echo "$job_entry" | cut -d ':' -f 2)
        local INPUT_FILE=$(echo "$job_entry" | cut -d ':' -f 3)
        local OUTPUT_LOG="$TEMP_OUTPUT_DIR/result_$job_i.log"
        
        if grep -q "Encryption time:" "$OUTPUT_LOG"; then
            SUCCESS=$((SUCCESS + 1))
        else
            echo "----------------------------------------------------------------------------------------------------------" >&2
            echo "【任务失败】Job $job_i (PID $pid) - 文件: $INPUT_FILE 未找到有效性能数据！" >&2
            if [ -f "$OUTPUT_LOG" ]; then
                echo "--- 测试工具输出/错误日志 ($OUTPUT_LOG) ---" >&2
                cat "$OUTPUT_LOG" >&2
            else
                echo "--- 错误: 找不到日志文件 $OUTPUT_LOG ---" >&2
            fi
            echo "----------------------------------------------------------------------------------------------------------" >&2
        fi
    done

    
    local ELAPSED_TIME_NS=$((END_TIME - START_TIME))
    local ELAPSED_TIME_S=$(awk "BEGIN {printf \"%.6f\", $ELAPSED_TIME_NS / 1000000000}")

    echo "[Jobs] 所有 $SUCCESS/$n 个进程完成。总耗时: ${ELAPSED_TIME_S} 秒。" >&2
    
    echo "$ELAPSED_TIME_S"

    if [ "$SUCCESS" -ne "$n" ]; then
        return 1
    fi
    return 0
}

# 功能: 统计测试结果并生成报告 (使用总文件大小和总耗时计算性能)
# 参数: $1 - 总文件数 (n)
# 参数: $2 - 总耗时 (秒)
aggregate_results() {
    local n=$1
    local total_elapsed_time_s=$2
    local total_enc_time_ns=0
    local total_dec_time_ns=0
    local total_files=0
    local total_file_size_bytes=0
    
    # 遍历所有日志文件
    for file in "$TEMP_OUTPUT_DIR/result_"*.log; do
        if [ -f "$file" ]; then
            
            if grep -q "Encryption time:" "$file"; then
                total_files=$((total_files + 1))
                
                # 从日志中提取时间和文件大小
                local enc_time=$(grep "Encryption time:" "$file" | awk '{print $3}')
                local file_size=$(grep "File size:" "$file" | awk '{print $3}')
                
                # 累加文件总大小
                if [ -n "$file_size" ]; then
                    total_file_size_bytes=$(awk "BEGIN {print $total_file_size_bytes + $file_size}")
                fi
                
                # 累加加密时间 (用于计算平均时间)
                if [ -n "$enc_time" ]; then
                    total_enc_time_ns=$(awk "BEGIN {print $total_enc_time_ns + $enc_time}")
                fi
                
                # 根据是否输出加密文件决定是否累加解密时间
                if [ "$OUTPUT_ENCRYPTED_FILES" = "false" ] && grep -q "Decryption time:" "$file"; then
                    local dec_time=$(grep "Decryption time:" "$file" | awk '{print $3}')
                    if [ -n "$dec_time" ]; then
                        total_dec_time_ns=$(awk "BEGIN {print $total_dec_time_ns + $dec_time}")
                    fi
                fi
            fi
        fi
    done
    
    if [ "$total_files" -eq 0 ]; then
        echo "错误: 未找到任何有效的测试结果文件进行聚合。" >&2
        return 1
    fi

    # 计算平均时间 (基于所有作业时间的总和)
    local avg_enc_time_sec=$(awk "BEGIN {printf \"%.6f\", $total_enc_time_ns / $total_files / 1000000000}")
    local avg_dec_time_sec=$(awk "BEGIN {printf \"%.6f\", $total_dec_time_ns / $total_files / 1000000000}")

    # 计算聚合性能 (Gbps) - 使用总文件大小和总耗时计算
    local total_file_size_bits=$(awk "BEGIN {print $total_file_size_bytes * 8}")
    
    # 加密性能 = 总文件大小(位) / 总耗时(秒)
    local agg_enc_gbps=0
    if [ "$(awk "BEGIN {print ($total_elapsed_time_s > 0) ? 1 : 0}")" -eq 1 ]; then
        agg_enc_gbps=$(awk "BEGIN {printf \"%.2f\", $total_file_size_bits / $total_elapsed_time_s / 1000000000}")
    fi
    
    # 解密性能 = 总文件大小(位) / 总耗时(秒)
    local agg_dec_gbps=0
    if [ "$(awk "BEGIN {print ($total_elapsed_time_s > 0) ? 1 : 0}")" -eq 1 ] && [ "$total_dec_time_ns" -gt 0 ]; then
        agg_dec_gbps=$(awk "BEGIN {printf \"%.2f\", $total_file_size_bits / $total_elapsed_time_s / 1000000000}")
    fi
    
    # 输出报告行
    printf "| %-10s | %-10s | %-18s | %-18s | %-18s | %-18s |\n" \
           "$MAX_CONCURRENCY" \
           "$total_files" \
           "$avg_enc_time_sec" \
           "$avg_dec_time_sec" \
           "$agg_enc_gbps" \
           "$agg_dec_gbps"
}

# --- 主执行逻辑 (已修改为单次运行) ---

echo ""
echo "=========================================================================================================="
echo "                           SM4 文件系统并发性能报告 (基于 $TEST_BIN)"
echo "=========================================================================================================="
if [ "$TEST_MODE" = "directory" ]; then
    echo "测试模式: 文件系统并发测试 ($INPUT_DIR -> $OUTPUT_DIR)"
else
    echo "测试模式: 内存模式测试 ($FILE_SIZE_STR x $FILE_COUNT -> $OUTPUT_DIR)"
fi
echo "算法/引擎: -elapsed -engine hct -cmd $CMD_TYPE -evp sm4"
echo "总核心数/NUMA节点: $N_CORES / $N_NODES"
echo "----------------------------------------------------------------------------------------------------------"
printf "| %-10s | %-10s | %-18s | %-18s | %-18s | %-18s |\n" "并发数 (n)" "完成文件" "平均加密时间(s)" "平均解密时间(s)" "聚合加密性能(Gbps)" "聚合解密性能(Gbps)"
printf "|%s|\n" "----------------------------------------------------------------------------------------------------------"

# 运行并发测试并获取总耗时
TOTAL_ELAPSED_TIME=$(run_concurrent_test)
if [ $? -eq 0 ]; then
    aggregate_results "$N_FILES" "$TOTAL_ELAPSED_TIME"
else
    printf "| %-10s | %-10s | %-18s | %-18s | %-18s | %-18s |\n" "$MAX_CONCURRENCY" "失败" "N/A" "N/A" "N/A" "N/A"
fi


echo "=========================================================================================================="

echo "测试完成。请查阅上面的性能报告。加密文件已输出到 $OUTPUT_DIR 目录。"

